package org.shx.springboot.ollama.ai.config;

import jakarta.annotation.PostConstruct;
import org.apache.http.util.Args;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.Generation;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.stereotype.Component;
import org.springframework.ai.chat.model.ChatResponse;

import reactor.core.publisher.Flux;

import java.lang.reflect.Method;
import java.util.List;
import java.util.stream.Collectors;

/**
 * 增强版Ollama日志切面：精准拦截OllamaChatModel，解决日志不输出问题
 */
@Aspect
@Component
public class OllamaLogAspect {

    private static final Logger logger = LoggerFactory.getLogger(OllamaLogAspect.class);

    // 初始化日志：确认切面加载
    @PostConstruct
    public void init() {
        logger.info("✅ OllamaLogAspect初始化成功，已准备拦截OllamaChatModel");
    }

    // ===== 拦截Controller方法（用于验证请求链路）=====
    @Around("execution(* org.shx.springboot.ollama.ai.controller.ChatController.*(..))")
    public Object logController(ProceedingJoinPoint joinPoint) throws Throwable {
        logger.info("\n===== 进入Controller方法: {} =====", joinPoint.getSignature().getName());
        Object result = joinPoint.proceed();
        logger.info("===== 离开Controller方法: {} =====", joinPoint.getSignature().getName());
        return result;
    }

    // ===== 核心：拦截OllamaChatModel方法（与Ollama直接交互的逻辑）=====
    // 切入点：OllamaChatModel类
    @Pointcut("within(org.springframework.ai.ollama.OllamaChatModel)")
    public void ollamaChatModelPointcut() {}

    /**
     * 拦截OllamaChatModel的stream方法（流式响应）
     * 解决：Flux响应拦截，打印实时返回内容
     */
    @Around("ollamaChatModelPointcut() && execution(* stream(..))")
    public Object logStream(ProceedingJoinPoint joinPoint) throws Throwable {

        // 直接从连接点获取参数（按顺序匹配）
        Object[] args = joinPoint.getArgs();
        Prompt prompt = (Prompt) args[0]; // 第一个参数是Prompt
        logger.info("\n▶️ 开始流式调用Ollama，输入提示: {}", getPromptPreview(prompt));

        // 执行目标方法，获取流式响应
        Flux<ChatResponse> responseFlux = (Flux<ChatResponse>) joinPoint.proceed();

        // 关键：通过反射调用ollamaChatRequest，手动获取请求参数（绕过私有方法拦截限制）
        try {
            // 获取OllamaChatModel的ollamaChatRequest方法（私有）
            // 1. 获取私有方法（指定参数类型：Prompt.class 和 boolean.class）
            Method chatRequestMethod = OllamaChatModel.class.getDeclaredMethod(
                    "ollamaChatRequest",
                    Prompt.class,  // 第一个参数类型
                    boolean.class  // 第二个参数类型
            );
            chatRequestMethod.setAccessible(true); // 强制访问私有方法
            // 调用方法获取请求对象（stream参数为true）
            OllamaApi.ChatRequest chatRequest = (OllamaApi.ChatRequest) chatRequestMethod.invoke(joinPoint.getTarget(), prompt, true);
            // 打印请求参数（解决400错误的核心信息）
            logger.info("\n📤 反射获取Ollama请求参数");
            // 打印请求详情（INFO级别，确保输出）
            logger.info("\n📤 发送请求到Ollama");
            logger.info("模型: {}", chatRequest.model());
            logger.info("messages: {}", chatRequest.messages());
            logger.info("options: {}", chatRequest.options());
            logger.info("是否流式: {}", chatRequest.stream());
            logger.info("系统提示: {}", extractSystemPrompt(prompt));
            logger.info("用户消息: {}", prompt.getInstructions().stream()
                    .filter(msg -> "user".equals(msg.getMessageType()))
                    .map(Message::getText)
                    .collect(Collectors.joining("; ")));
            logger.info("工具调用参数: {}", chatRequest.tools() != null ?
                    chatRequest.tools().stream().map(t -> t.function().name()).collect(Collectors.toList()) : "无");
        } catch (Exception e) {
            logger.warn("反射获取请求参数失败", e);
        }

        // 拦截流式响应
        return responseFlux.doOnNext(chunk -> {
            if (!chunk.getResults().isEmpty()) {
                String content = chunk.getResults().get(0).getOutput().getText();
                logger.info("📥 流式响应: {}", content.length() > 100 ? content.substring(0, 100) + "..." : content);
            }
        });
//        logger.info("\n▶️ 开始流式调用Ollama，输入提示: {}", getPromptPreview(prompt));
//        // 执行目标方法，获取流式响应
//        Flux<ChatResponse> responseFlux = (Flux<ChatResponse>) joinPoint.proceed();
//
//        // 拦截每个流式响应块（INFO级别打印关键内容）
//        return responseFlux.doOnNext(chunk -> {
//            if (!chunk.getResults().isEmpty()) {
//                Generation generation = chunk.getResults().get(0);
//                String content = generation.getOutput().getText();
//                logger.info("📥 流式响应: {}", content.length() > 100 ? content.substring(0, 100) + "..." : content);
//                // 打印元数据（判断是否完成）
//                logger.debug("是否完成: {}", chunk.getMetadata().get("done").toString());
//            }
//        });
    }

    /**
     * 拦截OllamaChatModel的call方法（同步响应）
     */
    @Around("ollamaChatModelPointcut() && execution(* call(..))")
    public Object logCall(ProceedingJoinPoint joinPoint) throws Throwable {
        // 直接从连接点获取参数（按顺序匹配）
        Object[] args = joinPoint.getArgs();
        Prompt prompt = (Prompt) args[0]; // 第一个参数是Prompt
        logger.info("\n▶️ 开始流式调用Ollama，输入提示: {}", getPromptPreview(prompt));
        // 执行目标方法，获取响应
        ChatResponse response = (ChatResponse) joinPoint.proceed();

        // 打印同步响应结果
        if (!response.getResults().isEmpty()) {
            logger.info("📥 同步响应: {}", response.getResults().get(0).getOutput().getText());
        } else {
            logger.warn("📥 同步响应为空");
        }
        return response;
    }

    // ===== 辅助方法：提取日志所需信息 =====
    /** 提取提示信息预览（避免日志过长） */
    private String getPromptPreview(Prompt prompt) {
        return prompt.getInstructions().stream()
                .map(msg -> msg.getMessageType() + ": " + (msg.getText().length() > 50 ? msg.getText().substring(0, 50) + "..." : msg.getText()))
                .collect(Collectors.joining("; "));
    }
    /** 提取系统提示 */
    private String extractSystemPrompt(Prompt prompt) {
        return prompt.getInstructions().stream()
                .filter(msg -> "system".equals(msg.getMessageType()))
                .map(msg -> msg.getText().length() > 100 ? msg.getText().substring(0, 100) + "..." : msg.getText())
                .findFirst().orElse("无系统提示");
    }

    /** 提取消息预览 */
    private List<String> extractMessagePreviews(List<OllamaApi.Message> messages) {
        return messages.stream()
                .map(msg -> msg.role() + ": " + (msg.content().length() > 50 ? msg.content().substring(0, 50) + "..." : msg.content()))
                .collect(Collectors.toList());
    }
}